import numpy as np
from scipy.signal import find_peaks, resample
from sklearn.cluster import DBSCAN
from butterworth_bandpass_filter import butterworth_bandpass_filter
from schmidt_spike_removal import schmidt_spike_removal
from energy_based_spike_removal import energy_based_spike_removal
import warnings
import scipy.io.wavfile as wav

warnings.filterwarnings('ignore', '.*Chunk (non-data) not understood, skipping it.*')

def check_heart_sound_validity(signal, Fs, mode='relaxed'):
    """
    检查心音信号的有效性。

    参数:
    - signal: 输入的心音信号
    - Fs: 采样频率
    - mode: 检测模式 ('relaxed' 或 'strict')

    返回:
    - is_valid: 是否有效
    - validTimeProportion: 有效时间比例
    - unique_clusters: 唯一的簇
    - clusterMedians: 每个簇的中位数
    - clusterFrequencies: 每个簇的频率
    - clusterTypes: 每个簇的类型
    """

    if mode == 'strict':
        # 严谨模式参数设置
        validityThreshold = 0.82  # 有效时间比例的阈值
        frequencyThreshold = 0.85  # 频域能量比例的阈值
        dutyCycleThresholdLow = 0.59  # 占空比的低阈值
        dutyCycleThresholdHigh = 0.64  # 占空比的高阈值
        zcrLowerThreshold = 80  # 零交叉率的下限阈值
        zcrUpperThreshold = 220  # 零交叉率的上限阈值
    else:
        # 宽松模式参数设置（默认）
        validityThreshold = 0.8  # 有效时间比例的阈值
        frequencyThreshold = 0.78  # 频域能量比例的阈值
        dutyCycleThresholdLow = 0.58  # 占空比的低阈值
        dutyCycleThresholdHigh = 0.65  # 占空比的高阈值
        zcrLowerThreshold = 80  # 零交叉率的下限阈值
        zcrUpperThreshold = 260  # 零交叉率的上限阈值

    min_length_threshold = 30  # 最小信号长度阈值

    # 降采样处理
    if Fs > 1000:
        num_samples = int(len(signal) * 1000 / Fs)
        signal = resample(signal, num_samples)
        Fs = 1000

    # 检查降采样后的信号长度是否足够
    if len(signal) < min_length_threshold:
        print(f"信号太短，只有{len(signal)}长度，低于{min_length_threshold}")
        return False, 0, [], [], [], []

    # 信号滤波处理
    signal = butterworth_bandpass_filter(signal, 4, 25, 400, Fs)
    signal, _ = schmidt_spike_removal(signal, Fs)
    signal, _ = energy_based_spike_removal(signal, Fs)

    # 信号标准化
    max_signal = np.max(np.abs(signal))
    if max_signal == 0:
        return False, 0, [], [], [], []
    signal = signal / max_signal

    # 频域能量比例计算
    energy_ratio = compute_energy_ratio(signal, Fs, [30, 200])
    is_valid_freq = energy_ratio >= frequencyThreshold

    # 零交叉率计算
    zcr_per_sec = compute_zcr_per_sec(signal, Fs)
    is_valid_zcr = zcrLowerThreshold <= zcr_per_sec <= zcrUpperThreshold

    # 占空比计算
    duty_cycle = compute_duty_cycle(signal)
    is_valid_duty = duty_cycle is not None and dutyCycleThresholdLow <= duty_cycle <= dutyCycleThresholdHigh

    # 峰值和间隔计算
    peaks, _ = find_peaks(signal, distance=0.15 * Fs, height=0.15)
    intervals = np.diff(peaks) / Fs

    if len(intervals) >= 2:
        db = DBSCAN(eps=0.1, min_samples=2).fit(intervals.reshape(-1, 1))
        labels = db.labels_
        unique_clusters = np.unique(labels).astype(float)
        clusterMedians = np.array([np.median(intervals[labels == x]) if len(intervals[labels == x]) > 0 else 0 for x in unique_clusters])
        clusterIntervals = np.array([np.sum(intervals[labels == x]) if len(intervals[labels == x]) > 0 else 0 for x in unique_clusters])
        total_interval_time = np.sum(intervals)
        clusterFrequencies = clusterIntervals / total_interval_time

        sortedIdx = np.argsort(clusterFrequencies)[::-1]
        if len(unique_clusters) > 2:
            for i in range(3, len(unique_clusters)):
                labels[labels == unique_clusters[sortedIdx[i]]] = -1
            unique_clusters = np.unique(labels).astype(float)
            clusterMedians = np.array([np.median(intervals[labels == x]) if len(intervals[labels == x]) > 0 else 0 for x in unique_clusters])
            clusterIntervals = np.array([np.sum(intervals[labels == x]) if len(intervals[labels == x]) > 0 else 0 for x in unique_clusters])
            clusterFrequencies = clusterIntervals / total_interval_time

        valid_clusters = (0.185 < clusterMedians) & (clusterMedians < 0.8)
        sum_of_valid_intervals = np.sum(clusterIntervals[valid_clusters])

        is_valid_time = (sum_of_valid_intervals / total_interval_time) >= validityThreshold
        validTimeProportion = sum_of_valid_intervals / total_interval_time
    else:
        is_valid_time = False
        validTimeProportion = 0
        unique_clusters = []
        clusterMedians = []
        clusterFrequencies = []
        valid_clusters = []

    is_valid = is_valid_freq and is_valid_zcr and is_valid_duty and is_valid_time
    clusterTypes = np.array([conditional_type(x, valid_clusters, unique_clusters) for x in unique_clusters])

    return is_valid, validTimeProportion, unique_clusters, clusterMedians, clusterFrequencies, clusterTypes

def compute_energy_ratio(signal, Fs, freq_range):
    if len(signal) == 0:
        return 0
    nfft = 2 ** int(np.ceil(np.log2(len(signal))))
    Y = np.fft.fft(signal, nfft) / len(signal)
    f = Fs / 2 * np.linspace(0, 1, nfft // 2 + 1)
    total_energy = np.sum(np.abs(Y[:nfft // 2 + 1]) ** 2)
    freq_idx = (f >= freq_range[0]) & (f <= freq_range[1])
    range_energy = np.sum(np.abs(Y[:nfft // 2 + 1][freq_idx]) ** 2)
    return range_energy / total_energy

def compute_zcr_per_sec(signal, Fs):
    if len(signal) == 0:
        return 0
    zero_crossings = np.sum(np.abs(np.diff(np.sign(signal)))) / 2
    duration = len(signal) / Fs
    return zero_crossings / duration

def compute_duty_cycle(signal):
    if len(signal) == 0:
        return None
    sorted_signal = np.sort(np.abs(signal))
    if len(sorted_signal) == 0:
        return None
    lower_30_percent = sorted_signal[:int(0.3 * len(sorted_signal))]
    if len(lower_30_percent) == 0:
        return None
    threshold = np.median(lower_30_percent) + 3 * np.std(lower_30_percent)
    return np.sum(np.abs(signal) > threshold) / len(signal)

def conditional_type(cluster_id, valid_clusters, unique_clusters):
    if cluster_id == -1:
        return 0  # 噪声
    else:
        idx = np.where(unique_clusters == cluster_id)[0][0]
        return 1 if valid_clusters[idx] else 2  # 有效或无效
